package priv.pfz.basedao.generator;

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.statement.SQLColumnDefinition;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlCreateTableStatement;
import com.alibaba.druid.util.JdbcConstants;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import priv.pfz.basedao.Utils;
import priv.pfz.basedao.apt.AptUtils;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.sql.*;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.*;

/**
 * 连接数据库，根据DDL自动生成Entity类
 * @author yyjxp
 * 2024/2/5 23:15
 */
public class Generator {

    private static final Map<String, FieldType> COLUMN_TYPE_MAP = ImmutableMap.<String, FieldType>builder()
            .put("char", FieldType.STRING)
            .put("varchar", FieldType.STRING)
            .put("text", FieldType.STRING)
            .put("enum", FieldType.STRING)
            .put("set", FieldType.STRING)
            .put("tinyint", FieldType.INTEGER)
            .put("smallint", FieldType.INTEGER)
            .put("mediumint", FieldType.INTEGER)
            .put("int", FieldType.LONG)
            .put("bigint", FieldType.LONG)
            .put("float", FieldType.BIG_DECIMAL)
            .put("double", FieldType.BIG_DECIMAL)
            .put("decimal", FieldType.BIG_DECIMAL)
            .put("datetime", FieldType.DATE)
            .put("date", FieldType.DATE)
            .put("time", FieldType.DATE)
            .put("timestamp", FieldType.DATE)
            .put("year", FieldType.DATE)
            .put("binary", FieldType.BYTE_ARRAY)
            .put("varbinary", FieldType.BYTE_ARRAY)
            .put("blob", FieldType.BYTE_ARRAY)
            .build();

    private static final DateTimeFormatter DTF = DateTimeFormatter.ofPattern("yyyy/M/d H:m");

    /**
     * entity类名的后缀
     * 默认为Entity
     * 例如user表生成UserEntity.java
     */
    private String entitySuffix = "Entity";

    /**
     * 存放entity的路径
     * 提示：可填写绝对路径或相对路径
     *      若填写相对路径，是基于System.getProperty("user.dir")
     *      一般Main方法的user.dir是在项目根目录，Junit是在Module根目录
     */
    private String entityPath;

    /**
     * entity的包名
     * 例如priv.pfz.basedao.lab.dao.entity
     */
    private String entityPackage;

    /**
     * 表所在数据库
     * 若dbUrl中选择了则无需再设置
     */
    private String dbName;

    /**
     * 需要生成entity的表名列表
     */
    private Set<String> tableNames = Sets.newLinkedHashSet();

    /**
     * 表名前缀，前缀不带入entity类名中
     * 例如tb_user表生成UserEntity.java
     */
    private String tableNamePrefix;

    /**
     * 表名后缀，后缀不带入entity类名中
     * 例如user_00表生成UserEntity.java
     */
    private String tableNameSuffix;

    /**
     * 数据库连接URL
     * 例如jdbc:mysql://host:port/databaseName
     */
    private String dbUrl;

    /**
     * 数据库用户名
     */
    private String dbUser;

    /**
     * 数据库密码
     */
    private String dbPassword;

    public static Generator newInstance() {
        return new Generator();
    }

    public Generator entitySuffix(String entitySuffix) {
        this.entitySuffix = entitySuffix;
        return this;
    }

    public Generator entityPath(String entityPath) {
        this.entityPath = entityPath;
        return this;
    }

    public Generator dbName(String dbName) {
        this.dbName = dbName;
        return this;
    }

    public Generator entityPackage(String entityPackage) {
        this.entityPackage = entityPackage;
        return this;
    }

    public Generator tableNames(String... tableNames) {
        Collections.addAll(this.tableNames, tableNames);
        return this;
    }

    public Generator tableNames(List<String> tableNames) {
        this.tableNames.addAll(tableNames);
        return this;
    }

    public Generator tableNamePrefix(String tableNamePrefix) {
        this.tableNamePrefix = tableNamePrefix;
        return this;
    }

    public Generator tableNameSuffix(String tableNameSuffix) {
        this.tableNameSuffix = tableNameSuffix;
        return this;
    }

    public Generator dbUrl(String dbUrl) {
        this.dbUrl = dbUrl;
        return this;
    }

    public Generator dbUser(String dbUser) {
        this.dbUser = dbUser;
        return this;
    }

    public Generator dbPassword(String dbPassword) {
        this.dbPassword = dbPassword;
        return this;
    }

    public void generate() {
        List<String> ddls = getDdls();
        for (String ddl : ddls) {
            generateOne(ddl);
        }
    }

    private List<String> getDdls() {
        // 加载并注册JDBC驱动类
        try {
            Class.forName("com.mysql.cj.jdbc.Driver");
        } catch (ClassNotFoundException e) {
            throw new RuntimeException("加载JDBC驱动异常", e);
        }

        List<String> ddls = Lists.newArrayList();
        // 建立数据库连接
        try (Connection conn = DriverManager.getConnection(dbUrl, dbUser, dbPassword)) {
            for (String tableName : tableNames) {
                Statement stmt = conn.createStatement();
                ResultSet rs = stmt.executeQuery(getQueryDdlSql(tableName));
                if (rs.next()) {
                    String createTableSQL = rs.getString(2);
                    System.out.println(createTableSQL);
                    ddls.add(createTableSQL);
                }
            }
        } catch (SQLException e) {
            throw new RuntimeException("查询DDL异常", e);
        }
        return ddls;
    }

    private String getQueryDdlSql(String tableName) {
        if (StringUtils.isBlank(dbName)) {
            return "SHOW CREATE TABLE `" + tableName + "`";
        } else {
            return "SHOW CREATE TABLE " + dbName + ".`" + tableName + "`";
        }
    }

    private void generateOne(String sql) {
        //调用druid解析DDL
        List<SQLStatement> stmtList = SQLUtils.parseStatements(sql, JdbcConstants.MYSQL);
        MySqlCreateTableStatement statement = (MySqlCreateTableStatement) stmtList.get(0);

        //tableName -> ClassName+Entity
        String tableName = removeShell(statement.getTableName());
        String entityName = tableToClass(tableName);

        String entityTemplate = readResource("entity_template");
        String fieldTemplate = readResource("field_template");
        Set<String> imports = Sets.newLinkedHashSet();
        List<String> fields = Lists.newArrayList();

        //column
        for (SQLColumnDefinition column : statement.getColumnDefinitions()) {
            //columnName -> fieldName
            String columnName = removeShell(column.getColumnName());
            String fieldName = AptUtils.columnToField(columnName);

            //columnType -> fieldType
            String columnType = column.getDataType().getName();
            FieldType fieldType = COLUMN_TYPE_MAP.get(columnType);
            if (fieldType == null) {
                throw new RuntimeException("不支持的数据库字段类型:" + columnType);
            }

            //columnComment -> fieldComment
            String comment = column.getComment() == null ? "" : removeShell(column.getComment().toString());

            Map<String, String> params = new HashMap<>();
            params.put("comment", comment);
            params.put("type", fieldType.getType().getSimpleName());
            params.put("name", fieldName);
            if (StringUtils.isNotBlank(fieldType.getImportStr())) {
                imports.add("import " + fieldType.getImportStr() + ";");
            }
            fields.add(AptUtils.fillTemplate(fieldTemplate, params));
        }

        Map<String, String> params = new HashMap<>();
        params.put("package", entityPackage);
        params.put("imports", String.join("\n", imports));
        params.put("date", LocalDateTime.now().format(DTF));
        params.put("classname", entityName);
        params.put("fields", String.join("\n", fields));
        String entity = AptUtils.fillTemplate(entityTemplate, params);
        String filename = Utils.getFilename(entityPath, entityName + ".java");
        writeSourceFile(filename, entity);
        System.out.println(filename + "生成完毕");
    }

    private String removeShell(String str) {
        return str.substring(1, str.length() - 1);
    }

    private String tableToClass(String tableName) {
        if (StringUtils.isNotBlank(tableNamePrefix) && tableName.startsWith(tableNamePrefix)) {
            tableName = tableName.substring(tableNamePrefix.length());
        }
        if (StringUtils.isNotBlank(tableNameSuffix) && tableName.endsWith(tableNameSuffix)) {
            tableName = tableName.substring(0, tableName.length() - tableNameSuffix.length());
        }
        return AptUtils.tableToClass(tableName) + entitySuffix;
    }

    private String readResource(String filename) {
        try (InputStream is = ClassLoader.getSystemResourceAsStream(filename)) {
            return IOUtils.toString(is, StandardCharsets.UTF_8);
        } catch (IOException e) {
            throw new RuntimeException("读取资源文件异常:" + filename, e);
        }
    }

    private void writeSourceFile(String filename, String content) {
        try (OutputStream os = Files.newOutputStream(Paths.get(filename))) {
            IOUtils.write(content, os, StandardCharsets.UTF_8);
        } catch (IOException e) {
            throw new RuntimeException("写java源文件异常:" + filename, e);
        }
    }
}
